#include "rec.h"

extern output op;  // global variable

//	@2022-02-12: add Helium .

// ============================================================================
// Notice that all physical quantities are expressed in unit of time, i.e., second.

int ODE_rec( double lna, const double y[], double f[], void *params ){

    cosmo_params *cp = (cosmo_params*)params;

    double a        = exp(lna);

    a = (a > 1.0) ? 1.0 : a;

    double Xe = y[0];
    
    /* extract cosmological parameters */

    double h        = cp->h;        // dimensionless ~70
    double Tcmb     = cp->Tcmb;     // 2.725K == 2.275*1.30909E11 s^{-1}
    double Omega_b  = cp->Omega_b;
    double Omega_r  = cp->Omega_r;
    double Omega_c  = cp->Omega_c;
    double Omega_L  = cp->Omega_L;

    // double ns       = cp->ns;
    double Yp       = cp->Yp;
    
    double Tb       = Tcmb/a;       // in uint K
    double Tb_eV    = Tb*kBoltz;    // in uint eV
    double Tb_s     = Tb*1.30909E11;    // in uint sec^{-1}
    
    double H0       = h * 3.24086E-18; // in sec^{-1}
    double H        = H0*sqrt( (Omega_b+Omega_c)/a/a/a 
                             + Omega_L
                             + Omega_r/a/a/a/a );
    
    // 1.8791 h^2 * 10^{−29} g cm^{−3}
    double rho_cr = 4.3124E50 * h*h;    // in uint sec^{-4}

    double nb = Omega_b*rho_cr/mH/a/a/a;    // in unit sec^{-3}
	double nH = (1.0-Yp)*nb;


    /* calculate the RHS of dXe/dx */

    //  phi2_Tb
    double phi2_Tb = 0.448*log(E0/Tb_eV);

    //  alpha2(Tb)
    double alpha2_Tb = (64*M_PI/sqrt(27*M_PI)) * pow(alpha/mE,2) * sqrt(E0/Tb_eV) * phi2_Tb;

    //  beta(Tb)
    double beta_Tb = alpha2_Tb * pow( mE*Tb_s/(2*M_PI), 1.5) * exp(-E0/Tb_eV);
    
    //  beta2(Tb)
    double beta2_Tb = beta_Tb*exp(0.75*E0/Tb_eV);
    
    //  n_1s
    double n_1s = (1-Xe)*nH;
    
    //  C_r(Tb)
    double Lambda_alpha = H*pow(3*E0,3)/64/M_PI/M_PI/n_1s;
    double C_r_Tb = (Lambda_2s_1s+Lambda_alpha) / (Lambda_2s_1s+Lambda_alpha+beta2_Tb);
    
    f[0] = C_r_Tb*(beta_Tb*(1-Xe)-nH*alpha2_Tb*Xe*Xe)/H;

    if( op.fp_xe_rhs != NULL ){
        fprintf(op.fp_xe_rhs, "%10.5f %10.8f %10.8f %15.10f\n", (1./a-1), lna, y[0], f[0]);
    }

    return GSL_SUCCESS;
}

void Saha_Xe( cosmo_params *cp, double lna, double *Xe ){

    double a = exp(lna);

    a = (a > 1.0) ? 1.0 : a;

    /* extract cosmological parameters */
    double h        = cp->h;        // dimensionless ~0.7
    double Tcmb     = cp->Tcmb;     // 2.725K == 2.275*1.30909E11 s^{-1}
    double Omega_b  = cp->Omega_b;

	double Yp       = cp->Yp;
    
    double Tb       = Tcmb/a;       // in uint K
    double Tb_eV    = Tb*kBoltz;    // in uint eV
    double Tb_s     = Tb*1.30909E11;    // in uint sec^{-1}
    
    // 1.8791 h^2 * 10^{−29} g cm^{−3}
    double rho_cr = 4.3124E50 * h*h;    // in uint sec^{-4}

    double nb = Omega_b*rho_cr/mH/a/a/a;    // in unit sec^{-3}

//    double nH = (1.0-Yp)*nb;

    //  solve Saha equation
	double fe = 1.0;    // start value
	double fe_next;     // next value of fe evaluated with fe

	int n_iter = 0;
	while(1){
		double ne = fe*nb;

		double mETb_2pi_3_2 = pow( mE*Tb_s/(2*M_PI), 1.5);
		double A1 = 2.0 * mETb_2pi_3_2 * exp(-XI0/Tb_eV) / ne;
		double A2 = 4.0 * mETb_2pi_3_2 * exp(-XI1/Tb_eV) / ne;
		double AH = mETb_2pi_3_2 * exp(-E0/Tb_eV) / ne;

///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//      NOTE: if you output nb, you will see numbers that are extremly large, but the finaly results are correct.
//              you can try using other unit(s) rather than 'second' adopted here.
/*		printf("nb = %15.8f ne = %15.8f A1 = %15.8f A2 = %15.8f AH = %15.8f\n", nb, ne, A1, A2, AH);*/
///////////////////////////////////////////////////////////////////////////////////////////////////////////////////

		fe_next = 0.25*Yp*(2.*A1*A2+A1)/(1+A1+A1*A2) + (1.-Yp)*AH/(1.+AH);

		fe = fe_next;
		if( fabs(fe_next-fe) < 1E-5 ){
			break;
		} 
		
		n_iter++;

		if( n_iter > 100 ){
			printf("Error: iterated more than 100 times, convergement failed!\n");
			exit(0);
		}
	}

	*Xe = fe/(1.0-Yp);
}

/*
    Xe_sw = 0.999
*/
void do_rec_calculation( cosmo_params *cp  ){

    printf("--> start solve the cosmic ionization history.\n");

    double z_start  = 8000;
    double a_ini    = 1./(1+z_start);
    double lna_ini  = log(a_ini);

    double z_stop   = 100;
    double a_stop   = 1./(1+z_stop);
    double lna_stop = log(a_stop);

    int steps = 10000;
    double dlna = (lna_stop-lna_ini)/(steps-1);

    double Xe;
    double ODE_Xe_ini;
    double ODE_lna_ini;

    int ode_switch_on = 0;

    if( op.fp_xe != NULL ){
        fprintf(op.fp_xe, "#   a       z       Xe\n");
    }

    // case 1: Xe is very close to 1.0
    int i;
    for( i=0; i<steps; i++ ){
        double lna = lna_ini + i*dlna;
        Saha_Xe( cp, lna, &Xe );

        if( op.fp_xe != NULL ){
            double a = exp(lna);
            double z = 1./a-1; 
            fprintf(op.fp_xe, "%15.12f %15.12f %15.12f\n", a, z, Xe);
        }
        
        if( Xe < XE_SW ){
            ode_switch_on   = 1;
            ODE_Xe_ini      = Xe;
            ODE_lna_ini     = lna;

            // printf("# finished solving Saha equation:\n");
            // printf("# ODE_lna_ini   = %15.10f\n", ODE_lna_ini);
            // printf("# ODE_Xe_ini    = %15.10f\n", ODE_Xe_ini);
            
            printf("--> finished Saha eqn. evaluations.\n");
            break;
        }
    }

    if( ode_switch_on == 1 ){
    
        printf("--> switch to solving recombination ODE.\n");
    
        const gsl_odeiv2_step_type  *T = gsl_odeiv2_step_rk8pd;
        gsl_odeiv2_step             *s = gsl_odeiv2_step_alloc(T, 1);
        gsl_odeiv2_control          *c = gsl_odeiv2_control_y_new(1e-10,0.0);
        gsl_odeiv2_evolve           *e = gsl_odeiv2_evolve_alloc(1);
        gsl_odeiv2_system           sys= { ODE_rec, NULL, 1, cp};

        double t    = ODE_lna_ini;
        double tmax = lna_stop;
        double h    = 1e-6;
        double y[1] = {ODE_Xe_ini};

        while( t < tmax ){

            int status = gsl_odeiv2_evolve_apply(e, c, s, &sys, &t, tmax, &h, y);

            if( status != GSL_SUCCESS )
                break;

            if( gsl_isnan(y[0]) ){
                break;
            }

            double a = exp(t);
            double z = 1./a-1;
            
            if( op.fp_xe != NULL ){
                fprintf(op.fp_xe, "%15.12f %15.12f %15.12f\n", a, z, y[0]);
                printf("%15.12f %15.12f %15.12f\n", a, z, y[0]);
            }
        }

        gsl_odeiv2_evolve_free(e);
        gsl_odeiv2_control_free(c);
        gsl_odeiv2_step_free(s);
        
        printf("--> finished solving recombination ODE.\n");
    }
}
